package com.tyros.common.secure.symmetrics;

import com.tyros.common.secure.CryptoException;

import javax.crypto.Cipher;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.security.Provider;
import java.security.SecureRandom;
import java.util.Arrays;

/**
 * AES_GCM_NoPadding
 *
 * @author zsy
 */
public class AesGcm implements Symmetric {

    private static final String ALGORITHM = "AES/GCM/NoPadding";

    private static final int TAG_LENGTH_BIT = 128;

    private static final int IV_LENGTH_BYTE = 12;

    private static final int KEY_UNLESS_BIT = 16;

    private final SecureRandom secureRandom;

    private final ThreadLocal<Cipher> cipherWrapper = new ThreadLocal<>();

    public AesGcm() {
        this(new SecureRandom());
    }

    public AesGcm(SecureRandom secureRandom) {
        this.secureRandom = secureRandom;
    }

    private Cipher getCipher(Provider provider) {
        Cipher cipher = cipherWrapper.get();
        if (cipher == null || cipher.getProvider() != provider) {
            try {
                if (provider != null) {
                    cipher = Cipher.getInstance(ALGORITHM, provider);
                } else {
                    cipher = Cipher.getInstance(ALGORITHM);
                }
            } catch (Exception e) {
                throw new CryptoException("Couldn't get cipher instance", e);
            }
            cipherWrapper.set(cipher);
            return cipherWrapper.get();
        } else {
            return cipher;
        }
    }

    @Override
    public byte[] decrypt(byte[] cipherData, byte[] password, byte[] additional, byte[] iv, Provider provider) {
        if (password.length < KEY_UNLESS_BIT) {
            throw new CryptoException("key length must be longer than 16 bytes");
        }
        try {
            final Cipher cipher = getCipher(provider);
            cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(password, "AES"),
                    new GCMParameterSpec(TAG_LENGTH_BIT, cipherData, 0, IV_LENGTH_BYTE));
            if (additional != null) {
                cipher.updateAAD(additional);
            }
            return cipher.doFinal(cipherData, IV_LENGTH_BYTE, cipherData.length - IV_LENGTH_BYTE);
        } catch (Exception e) {
            throw new CryptoException("Couldn't decrypt", e);
        }
    }

    @Override
    public byte[] encrypt(byte[] raw, byte[] srcPassword, byte[] additional, byte[] iv, Provider provider) {
        if (srcPassword.length < KEY_UNLESS_BIT) {
            throw new CryptoException("key length must be longer than 16 bytes");
        }
        if (iv != null && iv.length != IV_LENGTH_BYTE) {
            throw new CryptoException(
                    "Expected salt value is " + IV_LENGTH_BYTE + ",but current value is " + iv.length);
        }
        byte[] encrypted = null;
        try {
            if (iv == null) {
                iv = new byte[IV_LENGTH_BYTE];
                synchronized (secureRandom) {
                    secureRandom.nextBytes(iv);
                }
            }
            final Cipher cipher = getCipher(provider);
            cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(srcPassword, "AES"),
                    new GCMParameterSpec(TAG_LENGTH_BIT, iv));
            if (additional != null) {
                cipher.updateAAD(additional);
            }
            encrypted = cipher.doFinal(raw);
            ByteBuffer byteBuffer = ByteBuffer.allocate(IV_LENGTH_BYTE + encrypted.length);
            byteBuffer.put(iv);
            byteBuffer.put(encrypted);
            return byteBuffer.array();
        } catch (Exception e) {
            throw new CryptoException("Couldn't encrypt", e);
        } finally {
            clear(iv);
            clear(encrypted);
        }

    }

    /**
     * 清理内存
     *
     * @param bs
     */
    private void clear(byte[] bs) {
        if (bs != null) {
            Arrays.fill(bs, (byte) 0);
        }
    }
}
